# coding:utf-8
# __author__ = yuan
# __time__ = 2020/3/16
# __file__ = net
# __desc__ =
# coding:utf-8
# __author__ = yuan
# __time__ = 2020/3/16
# __file__ = Backbone
# __desc__ =
# coding:utf-8
# __user__ = hiicy redldw
# __time__ = 2020/1/14
# __file__ = resnet
# __desc__ =
import torch.nn.functional as F
import torch
import torch.nn as nn
from torchvision.models import resnet101, resnet50


def get_resnet50_4():
    res = resnet50(pretrained=False)
    cres = list(res.children())[:5]
    res = nn.Sequential(*cres)
    return res


def convbn(inc, outc, kernel=1, stride=1, pad=0):
    return nn.Sequential(
        nn.Conv2d(inc, outc, kernel, stride, pad),
        nn.BatchNorm2d(outc),
    )


class BottleBlock(nn.Module):
    def __init__(self, inc, outc=256, downsample=None):
        super(BottleBlock, self).__init__()
        self.conv1 = convbn(inc, 128, 1)
        self.conv2 = convbn(128, 128, 3, pad=1)
        self.conv3 = convbn(128, outc, 1)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        identity = x
        x = self.conv1(x)
        x = self.relu(x)

        x = self.conv2(x)
        x = self.relu(x)

        out = self.conv3(x)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)
        return out


class Hourglass(nn.Module):
    def __init__(self, inc, outc, degree=4):
        super(Hourglass, self).__init__()
        self.inc = inc
        self.outc = outc
        rsis, pools = [], []
        self._make_hourglass(degree, rsis, pools, inc, outc)
        self.upresidule = []
        for _ in range(degree):
            self.upresidule.append(
                BottleBlock(256, 256)
            )
        self.rsis = rsis
        self.pools = pools
        del rsis
        del pools

    def _make_hourglass(self, degree, rsis: list, pools, inc, outc):
        if degree == 0:
            return
        residual = nn.Sequential(
            BottleBlock(inc, 256),
            BottleBlock(256, 256),
            BottleBlock(256, outc),  # 因为是256,所以我不用加下采样层
        )
        rsis.append(residual)
        pools.append(nn.Sequential(
            nn.MaxPool2d((1, 1), (2, 2)),
            BottleBlock(256, 256),
            BottleBlock(256, 256),
            BottleBlock(256, 256),
        ))
        self._make_hourglass(degree - 1, rsis, pools, 256, outc)

    def forward(self, x):
        """
        :param x:256通道
        """
        upori = [self.rsis[0](x)]  # 1个 # 4张上位图
        for i in range(0, len(self.rsis[1:])):
            downmap = self.pools[i](x)
            upori.append(self.rsis[i + 1](downmap))  # 1+1+1
        downmap = self.pools[-1](downmap)
        upori.reverse()  # 原地变换顺序
        for i in range(len(self.upresidule)):
            upmap = self.upresidule[i](downmap)
            uppermap = upori[i]
            size = uppermap.shape
            upmap = F.interpolate(upmap, size[2:], mode="bilinear", align_corners=True)
            upmap += uppermap

        return upmap


def conv3x3bnlu(inc, outc, kernel=3, stride=1, pad=0):
    from collections import OrderedDict
    return nn.Sequential(OrderedDict([
        ("conv", nn.Conv2d(inc, outc, kernel, stride, pad)),
        ("bn", nn.BatchNorm2d(outc)),
        ("relu", nn.ReLU(True))
    ]))


class WordDetection(nn.Module):
    def __init__(self, inc, newc):
        super(WordDetection, self).__init__()
        self.det_feature = conv3x3bnlu(inc, newc, 3)
        self.regression_fea = conv3x3bnlu(newc, newc, 3)
        self.fb_fea = conv3x3bnlu(newc, newc, 3)

        self.segmentation = nn.Conv2d(newc, 1, 1)
        self.regression = nn.Conv2d(newc, 5, 1)  # FIXME:可以八个通道或目前的5通道
        self.ori_orient = nn.Conv2d(newc, 1, 1)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):  # x: w/4, h/4
        fea = self.det_feature(x)
        reg_fea = self.regression_fea(fea)
        fb_fea = self.fb_fea(fea)

        segmentation = self.segmentation(fb_fea)
        segmentation = self.sigmoid(segmentation)  # 是或不是文本的概率
        regression = self.regression(reg_fea)
        orient = self.ori_orient(reg_fea)

        return (segmentation, regression, orient)


class CharDetection(nn.Module):
    def __init__(self, inc, outc):
        super(CharDetection, self).__init__()
        self.det_feature = conv3x3bnlu(inc, outc)
        self.reg_fea = conv3x3bnlu(outc, outc)
        self.fb_fea = conv3x3bnlu(outc, outc)

        self.sigmoid = nn.Sigmoid()

        self.segment = nn.Conv2d(outc, 1, 1)
        self.regression = nn.Conv2d(outc, 1, 1)
        self.ori_orient = nn.Conv2d(outc, 1, 1)

    def forward(self, x):
        fea = self.det_feature(x)
        reg_fea = self.reg_fea(fea)
        fb_fea = self.fb_fea(fea)

        segmentation = self.segment(fb_fea)
        segmentation = self.sigmoid(segmentation)  # 是或不是文本的概率
        regression = self.regression(reg_fea)
        orient = self.ori_orient(reg_fea)

        return (segmentation, regression, orient)

class CharRecognition(nn.Module):
    def __init__(self,inc,outc,numclass):
        super(CharRecognition, self).__init__()
        self.recog_fea = conv3x3bnlu(inc,outc)
        self.final_fea = conv3x3bnlu(outc,outc)

        self.recognition = nn.Conv2d(outc,numclass,1)
    def forward(self,x):
        fea = self.recog_fea(x)
        recog_fea = self.final_fea(fea)

        output = self.recognition(recog_fea)
        return output

class CharNet(nn.Module):
    def __init__(self, inc,outc, numclass):
        super(CharNet, self).__init__()
        self.body = get_resnet50_4()  # 基本模型50 4倍下采样

        self.hourglass1 = Hourglass(inc, inc)
        self.hourglass2 = Hourglass(inc, inc)

        self.wordbranch = WordDetection(inc,outc)
        self.charbranch_detection = CharDetection(inc,outc)
        self.charbranch_recogintion = CharRecognition(inc,outc,numclass)

    def forward(self, x):
        base_feature = self.body(x)
        hour_fea1 = self.hourglass1(base_feature)
        hour_fea2 = self.hourglass2(hour_fea1)

        word_seg, word_regression, word_orient = self.wordbranch(hour_fea2)
        char_seg, char_regression, char_orient = self.charbranch_detection(hour_fea2)
        char_cls = self.charbranch_recogintion(hour_fea2)

        return [[word_seg,word_regression,word_orient],
                [char_seg,char_regression,char_orient],
                [char_cls]]


if __name__ == '__main__':
    # hg = Hourglass(256, 256, 4)
    # p = hg(x)
    cn = CharNet(256,512,68)
    x = torch.rand(1, 3, 400, 400)
    p = cn(x)
    def show(p):
        if isinstance(p,list):
            for pp in p:
                if isinstance(pp,list):
                    show(pp)
                else:
                    print(pp.shape)
    show(p)